package uav.util.deepsort;

import org.ejml.simple.SimpleMatrix;

// 完整的卡尔曼滤波器实现
public class KalmanFilter {
    private SimpleMatrix state;      // [x, y, vx, vy]^T
    private SimpleMatrix covariance;
    private final SimpleMatrix F;    // 状态转移矩阵
    private final SimpleMatrix Q;    // 过程噪声
    private final SimpleMatrix H;    // 观测矩阵
    private final SimpleMatrix R;    // 观测噪声

    public KalmanFilter(float x, float y) {
        // 初始化状态向量（位置+速度）
        state = new SimpleMatrix(4, 1, true, new double[]{x, y, 0, 0});

        // 初始化协方差矩阵
        covariance = SimpleMatrix.identity(4).scale(1000);

        // 状态转移矩阵（匀速模型）
        F = new SimpleMatrix(4, 4, true, new double[]{
                1, 0, 1, 0,
                0, 1, 0, 1,
                0, 0, 1, 0,
                0, 0, 0, 1
        });

        // 过程噪声矩阵
        Q = SimpleMatrix.identity(4).scale(1e-4);

        // 观测矩阵（只能观测位置）
        H = new SimpleMatrix(2, 4, true, new double[]{
                1, 0, 0, 0,
                0, 1, 0, 0
        });

        // 观测噪声
        R = SimpleMatrix.identity(2).scale(0.1);
    }

    public void predict() {
        // 状态预测
        state = F.mult(state);
        // 协方差预测
        covariance = F.mult(covariance).mult(F.transpose()).plus(Q);
    }

    public void update(double[] measurement) {
        SimpleMatrix z = new SimpleMatrix(2, 1, true, new double[]{measurement[0],measurement[1]});

        // 计算卡尔曼增益
        SimpleMatrix S = H.mult(covariance).mult(H.transpose()).plus(R);
        SimpleMatrix K = covariance.mult(H.transpose()).mult(S.invert());

        // 状态更新
        state = state.plus(K.mult(z.minus(H.mult(state))));

        // 协方差更新
        covariance = covariance.minus(K.mult(H).mult(covariance));
    }

    public double[] getPosition() {
        return new double[]{state.get(0), state.get(1)};
    }
}
